import pandas as pd
import requests
from sqlalchemy.orm import Session

from app.bigdata.models.bdl_disk_capacity_info import BdlDiskCapacityInfo
from core.logger import logger


def query_prometheus_data(prometheus_url, expression):
    try:
        url = f"{prometheus_url}api/v1/query?query={expression}"
        req = requests.get(url=url)
        if req.status_code != 200:
            logger.error(f"query prometheus expression:{expression} failed => {req.text}")
            return False

    except Exception as e:
        logger.error(f"query prometheus expression:{expression} error => {e}")
        return False

    logger.info(f"query prometheus expression:{expression} success")
    return req.json()['data']['result']


def bdl_disk_capacity_handler():
    data = {
        "worker01": "192.168.63.182",
        "worker02": "192.168.63.185",
        "worker03": "192.168.63.183",
        "worker04": "192.168.63.184",
        "worker05": "192.168.63.186",
        "worker06": "192.168.63.191",
        "worker07": "192.168.63.189",
        "worker08": "192.168.63.188",
        "worker09": "192.168.63.190",
        "worker10": "192.168.63.197",
        "worker11": "192.168.63.198",
        "worker12": "192.168.63.199",
        "worker13": "192.168.63.200",
        "worker14": "192.168.63.201",
        "worker15": "192.168.63.202",
        "worker22": "192.168.63.209",
        "worker23": "192.168.63.221",
        "worker24": "192.168.63.222",
        "worker25": "192.168.63.223",
        "worker33": "192.168.63.246",
        "worker34": "192.168.63.247",
        "worker35": "192.168.63.248",
        "worker36": "192.168.63.249",
        "worker37": "192.168.63.250"
    }
    result = []
    for k, v in data.items():
        disk_total_expr = "disk_total{{ident=\"{0}-{1}\",path=~\"/disk1|/disk3|/disk4|/disk5\"}}".format(k, v)
        disk_free_expr = "disk_free{{ident=\"{0}-{1}\",path=~\"/disk1|/disk3|/disk4|/disk5\"}}".format(k, v)
        disk_total_data = query_prometheus_data("http://192.168.43.25:25526/", disk_total_expr)
        disk_free_data = query_prometheus_data("http://192.168.43.25:25526/", disk_free_expr)

        if disk_total_data and disk_free_data:
            for i in disk_total_data:
                for j in disk_free_data:
                    if i['metric']['path'] == j['metric']['path']:
                        disk_total = round(int(i['value'][1]) / 1024 / 1024 / 1024 / 1024, 2)
                        disk_free = round(int(j['value'][1]) / 1024 / 1024 / 1024 / 1024, 2)
                        result.append({
                            "ident": i['metric']['ident'],
                            "path": i['metric']['path'],
                            "disk_total": disk_total,
                            "disk_free": disk_free,
                            "usage_rate": (1 - round(disk_free / disk_total, 2)) * 100
                        })
    return result


def disk_hdfs_capacity_report_handler(db: Session, start_time: str, end_time: str):
    try:
        result = BdlDiskCapacityInfo.get_hdfs_total_capacity(db, start_time, end_time)
        if not result:
            return [], []
        df = pd.DataFrame(result)
        df = df.T
        df.reset_index(drop=False, inplace=True)
        df.rename(columns={"index": "容量状态"}, inplace=True)
        # 找到 'create_date' 所在的行
        create_date_row = df[df['容量状态'] == 'create_date']
        # 使用 'create_date' 的值来替换列名
        df.columns = ['容量状态'] + list(create_date_row.iloc[0, 1:])
        # 删除 'create_date' 这一行
        df = df[df['容量状态'] != 'create_date']
        data = df.to_dict(orient='split')
        title, content = data['columns'], data['data']
        return title, content
    except Exception as e:
        logger.error(f"disk_capacity_report_handler => {e}")
        return [], []


def disk_kudu_capacity_report_handler(db: Session, start_time: str, end_time: str):
    try:
        result = BdlDiskCapacityInfo.get_kudu_total_capacity(db, start_time, end_time)
        if not result:
            return [], []
        df = pd.DataFrame(result)
        df = df.T
        df.reset_index(drop=False, inplace=True)
        df.rename(columns={"index": "容量状态"}, inplace=True)
        # 找到 'create_date' 所在的行
        create_date_row = df[df['容量状态'] == 'create_date']
        # 使用 'create_date' 的值来替换列名
        df.columns = ['容量状态'] + list(create_date_row.iloc[0, 1:])
        # 删除 'create_date' 这一行
        df = df[df['容量状态'] != 'create_date']
        data = df.to_dict(orient='split')
        title, content = data['columns'], data['data']
        return title, content
    except Exception as e:
        logger.error(f"disk_capacity_report_handler => {e}")
        return [], []


if __name__ == "__main__":
    from common.database import SessionLocal

    session = SessionLocal()

    disk_hdfs_capacity_report_handler(session, "2024-01-29", "2024-02-01")
    disk_kudu_capacity_report_handler(session, "2024-01-29", "2024-02-01")
